import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
import cv2
import numpy as np
import serial
import time
import math

# -------------------------- 协议常量定义（与STM32匹配） --------------------------
FRAME_HEADER = 0x66  # 帧头
FRAME_LENGTH = 0x07  # 帧长度（固定7字节）
FRAME_TAIL = 0x5B    # 帧尾
SEND_INTERVAL = 0.05  # 发送间隔（20Hz，避免数据拥堵）

class ImageProcessorNode(Node):
    def __init__(self):
        super().__init__('image_processor_node')
        
        # 初始化相机参数和组件
        self.cam_params = self.init_camera_params()
        self.stereo, self.right_matcher, self.wls_filter = self.init_stereo_matcher()
        self.ser = self.init_serial()
        self.screen_center = (self.cam_params["size"][0] // 2, self.cam_params["size"][1] // 2)  # 图像中心点
        self.get_logger().info(f"图像中心点: {self.screen_center}")
        
        # 创建CV桥接器
        self.bridge = CvBridge()
        
        # 创建图像订阅者
        self.image_sub = self.create_subscription(
            Image,
            'image_raw',  # usb_cam默认发布的话题名
            self.image_callback,
            10)
        self.image_sub  # 防止被垃圾回收
        
        # 创建处理后图像发布者（可选，用于调试）
        self.processed_image_pub = self.create_publisher(Image, 'processed_image', 10)
        
        # 初始化变量
        self.last_send_time = time.time()
        self.last_image = None
        
        self.get_logger().info("图像处理器节点已启动，等待图像数据...")

    def init_camera_params(self):
        """初始化双目相机参数（包含畸变矫正参数）"""
        # 保持与原代码相同的相机参数初始化
        left_camera_matrix = np.array([[715.088662856028, -1.27750597113329, 321.725929157012],
                                       [0, 717.938092436680, 251.352481762661],
                                       [0., 0., 1.]])
        right_camera_matrix = np.array([[712.537452801351, -1.49459241948369, 325.232106968860],
                                        [0, 715.552769066380, 250.629910089625],
                                        [0., 0., 1.]])

        left_distortion = np.array(
            [[-0.440986151974959, 0.327464329751247, -0.000547150225973162, -0.00126940700042557, -0.335074894313998]])
        right_distortion = np.array(
            [[-0.444122749137210, 0.348225581098782, 0.00176792559350367, 0.00123611244757822, -0.660993977721134]])

        R = np.array([[0.999866011965132, 0.000416637643297468, 0.0163641232584165],
                      [-0.000544744807817963, 0.999969236927380, 0.00782486114692518],
                      [-0.0163603597159969, -0.00783272698033651, 0.999835479975588]])
        T = np.array([-59.8387197927031, 0.0107452720401373, 0.330883446596301])
        size = (640, 480)

        # 计算立体校正参数
        R1, R2, P1, P2, Q, validPixROI1, validPixROI2 = cv2.stereoRectify(
            left_camera_matrix, left_distortion,
            right_camera_matrix, right_distortion,
            size, R, T
        )

        # 生成左右相机的矫正映射表（用于后续图像矫正）
        left_map1, left_map2 = cv2.initUndistortRectifyMap(
            left_camera_matrix, left_distortion, R1, P1, size, cv2.CV_16SC2)
        right_map1, right_map2 = cv2.initUndistortRectifyMap(
            right_camera_matrix, right_distortion, R2, P2, size, cv2.CV_16SC2)

        return {
            "left_camera_matrix": left_camera_matrix,
            "right_camera_matrix": right_camera_matrix,
            "left_distortion": left_distortion,
            "right_distortion": right_distortion,
            "R": R,
            "T": T,
            "size": size,
            "Q": Q,
            "left_map1": left_map1,
            "left_map2": left_map2,
            "right_map1": right_map1,
            "right_map2": right_map2,
            "baseline": abs(T[0])
        }

    def init_stereo_matcher(self, blockSize=5, numDisparities=128):
        """初始化SGBM匹配器和WLS滤波器（用于视差计算）"""
        img_channels = 3
        stereo = cv2.StereoSGBM_create(
            minDisparity=0,
            numDisparities=numDisparities,
            blockSize=blockSize,
            P1=8 * img_channels * blockSize * blockSize,
            P2=32 * img_channels * blockSize * blockSize,
            disp12MaxDiff=5,
            preFilterCap=63,
            uniquenessRatio=15,
            speckleWindowSize=200,
            speckleRange=32,
            mode=cv2.STEREO_SGBM_MODE_SGBM_3WAY
        )

        right_matcher = cv2.ximgproc.createRightMatcher(stereo)
        wls_filter = cv2.ximgproc.createDisparityWLSFilter(matcher_left=stereo)
        wls_filter.setLambda(8000)
        wls_filter.setSigmaColor(1.5)

        return stereo, right_matcher, wls_filter

    def init_serial(self, port='/dev/ttyUSB0', baudrate=115200):
        """初始化串口通信（与STM32通信）"""
        try:
            ser = serial.Serial(
                port=port,
                baudrate=baudrate,
                bytesize=serial.EIGHTBITS,
                parity=serial.PARITY_NONE,
                stopbits=serial.STOPBITS_ONE,
                timeout=0.1
            )
            self.get_logger().info(f"串口已打开: {ser.port}")
            return ser
        except serial.SerialException as e:
            self.get_logger().error(f"串口打开失败: {e}")
            self.get_logger().info("将以无串口模式运行")
            return None

    def image_callback(self, msg):
        """处理接收到的图像消息"""
        try:
            # 将ROS图像消息转换为OpenCV格式
            cv_image = self.bridge.imgmsg_to_cv2(msg, "bgr8")
            
            # 分割左右相机图像（假设图像是1280x480，左右各640x480）
            frame_left = cv_image[:, :640]
            frame_right = cv_image[:, 640:]
            
            # 生成原始灰度图
            gray_left = cv2.cvtColor(frame_left, cv2.COLOR_BGR2GRAY)
            
            # 对左相机图像进行畸变矫正（彩色图和灰度图）
            left_rect_gray = self.rectify_image(gray_left, self.cam_params["left_map1"], self.cam_params["left_map2"])
            left_rect_color = self.rectify_image(frame_left, self.cam_params["left_map1"], self.cam_params["left_map2"])
            right_rect_gray = self.rectify_image(cv2.cvtColor(frame_right, cv2.COLOR_BGR2GRAY),
                                            self.cam_params["right_map1"], self.cam_params["right_map2"])
            
            # 计算视差和深度
            _, _, disp_color, threeD = self.compute_disparity_and_depth(
                left_rect_gray, right_rect_gray, self.stereo, self.right_matcher, self.wls_filter, self.cam_params["Q"]
            )
            
            # 在矫正后的彩色图上检测橙色球
            ball_info = self.detect_orange_ball(left_rect_color, threeD, self.cam_params["left_map1"])
            
            # 绘制检测结果（在矫正后的彩色图上）
            target_x, target_y = self.screen_center  # 默认发送中心点坐标
            if ball_info:
                # 获取检测信息
                cx_rect, cy_rect = ball_info["rect_center"]
                distance = ball_info["distance"]
                contour = ball_info["contour"]
                target_x, target_y = cx_rect, cy_rect  # 更新发送坐标为小球中心
                
                # 绘制小球轮廓
                cv2.drawContours(left_rect_color, [contour], -1, (0, 255, 0), 2)
                # 绘制小球中心
                cv2.circle(left_rect_color, (cx_rect, cy_rect), 5, (0, 0, 255), -1)
                # 绘制与中心点的连线
                cv2.line(left_rect_color, self.screen_center, (cx_rect, cy_rect), (0, 255, 255), 2)
                # 显示坐标和距离（基于彩色图圆心坐标）
                cv2.putText(left_rect_color, f"Center: ({cx_rect}, {cy_rect})", (10, 30),
                          cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 0), 2)
                cv2.putText(left_rect_color, f"Distance: {distance:.3f}m", (10, 60),
                          cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 0), 2)
            else:
                # 未检测到目标时显示提示
                cv2.putText(left_rect_color, "No Target", (10, 30),
                          cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
            
            # 绘制图像中心点标记
            cv2.circle(left_rect_color, self.screen_center, 5, (255, 0, 0), -1)
            cv2.line(left_rect_color, (self.screen_center[0]-10, self.screen_center[1]),
                    (self.screen_center[0]+10, self.screen_center[1]), (255, 0, 0), 2)
            cv2.line(left_rect_color, (self.screen_center[0], self.screen_center[1]-10),
                    (self.screen_center[0], self.screen_center[1]+10), (255, 0, 0), 2)
            
            # 按间隔发送坐标（20Hz）
            current_time = time.time()
            if current_time - self.last_send_time > SEND_INTERVAL:
                self.send_coordinate(target_x, target_y)
                self.last_send_time = current_time
            
            # 显示处理后的图像
            cv2.imshow("Processed Image", left_rect_color)
            cv2.waitKey(1)
            
            # 发布处理后的图像（可选）
            processed_msg = self.bridge.cv2_to_imgmsg(left_rect_color, "bgr8")
            self.processed_image_pub.publish(processed_msg)
            
        except Exception as e:
            self.get_logger().error(f"图像处理出错: {e}")

    def rectify_image(self, img, map1, map2):
        """对图像进行畸变矫正（支持彩色图和灰度图）"""
        return cv2.remap(img, map1, map2, cv2.INTER_LINEAR)

    def compute_disparity_and_depth(self, left_rect, right_rect, stereo, right_matcher, wls_filter, Q, depth_scale=16.0):
        """计算视差图和深度信息"""
        disparity_left = stereo.compute(left_rect, right_rect)
        disparity_right = right_matcher.compute(right_rect, left_rect)

        filtered_disparity = wls_filter.filter(disparity_left, left_rect, None, disparity_right)
        filtered_disparity = np.clip(filtered_disparity, 0, None)

        # 视差图归一化（用于显示）
        disp_norm = cv2.normalize(filtered_disparity, None, 0, 255, cv2.NORM_MINMAX, cv2.CV_8U)
        disp_color = cv2.applyColorMap(disp_norm, cv2.COLORMAP_JET)

        # 计算三维坐标（深度信息）
        threeD = cv2.reprojectImageTo3D(filtered_disparity, Q, handleMissingValues=False)
        threeD = threeD * depth_scale  # 缩放深度值

        return filtered_disparity, disp_norm, disp_color, threeD

    def detect_orange_ball(self, rect_color, threeD, left_map1):
        """在矫正后的彩色图上检测橙色球（确保坐标与深度图对齐）"""
        # 转换到HSV颜色空间进行颜色检测
        hsv = cv2.cvtColor(rect_color, cv2.COLOR_BGR2HSV)
        lower_orange = np.array([15, 82, 210])  # 橙色HSV范围（适配不同光照）
        upper_orange = np.array([68, 255, 255])
        mask = cv2.inRange(hsv, lower_orange, upper_orange)

        # 形态学操作去除噪声
        kernel = np.ones((5, 5), np.uint8)
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)  # 闭合孔洞
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)   # 去除小噪点

        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        ball_info = None
        best_score = 0

        for contour in contours:
            area = cv2.contourArea(contour)
            if area < 50:  # 过滤小轮廓
                continue

            # 计算圆形度（筛选圆形物体）
            perimeter = cv2.arcLength(contour, True)
            if perimeter == 0:
                continue
            circularity = 4 * math.pi * area / (perimeter **2)
            if circularity < 0.6:  # 只保留接近圆形的物体
                continue

            # 计算轮廓中心（矫正后坐标）
            (cx_rect, cy_rect), radius = cv2.minEnclosingCircle(contour)
            cx_rect, cy_rect = int(cx_rect), int(cy_rect)

            # 验证深度有效性
            if 0 <= cx_rect < threeD.shape[1] and 0 <= cy_rect < threeD.shape[0]:
                depth = threeD[cy_rect, cx_rect]
                z_val = depth[2]
                if z_val <= 0:
                    continue

                # 计算实际距离（米）
                distance = np.linalg.norm(depth) / 1000.0

                # 评分（优先选择面积大、圆形度高的目标）
                score = area * circularity
                if score > best_score:
                    best_score = score
                    ball_info = {
                        "rect_center": (cx_rect, cy_rect),
                        "depth": depth,
                        "distance": distance,
                        "contour": contour
                    }

        return ball_info

    def send_coordinate(self, x, y):
        """按协议发送坐标给下位机"""
        if not self.ser or not self.ser.is_open:
            return

        # 限制坐标范围为0~65535（16位无符号整数）
        x_clamped = np.clip(int(x), 0, 65535)
        y_clamped = np.clip(int(y), 0, 65535)

        # 拆分高低字节
        x_low = x_clamped & 0xFF
        x_high = (x_clamped >> 8) & 0xFF
        y_low = y_clamped & 0xFF
        y_high = (y_clamped >> 8) & 0xFF

        # 构建帧结构：[帧头, 长度, X低, X高, Y低, Y高, 帧尾]
        frame = bytearray([
            FRAME_HEADER, FRAME_LENGTH,
            x_low, x_high, y_low, y_high,
            FRAME_TAIL
        ])

        # 发送数据
        try:
            self.ser.write(frame)
            self.get_logger().info(f"发送坐标: ({x_clamped}, {y_clamped})  帧: {[hex(b) for b in frame]}")
        except serial.SerialException as e:
            self.get_logger().error(f"串口发送失败: {e}")

def main(args=None):
    rclpy.init(args=args)
    node = ImageProcessorNode()
    rclpy.spin(node)
    node.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()
